package com.fosss.imageRecognition;

import javax.imageio.ImageIO;
import javax.swing.*;
import javax.swing.filechooser.FileNameExtensionFilter;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.UUID;

import net.sourceforge.tess4j.ITessAPI;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;

public class ImageRecognitionUI extends JFrame {
    //语言包位置
    public static final String tessDataPath = "F:\\Tesseract-OCR\\tessdata";
    //截屏存储的位置
    public static final String screenshotPath = "F:\\Tesseract-OCR\\screenshot\\";

    private JButton uploadButton;  // 上传按钮
    private JLabel imageLabel;  // 图像标签
    private JTextArea textArea;  // 文本区域
    private JButton clearButton;  // 清空按钮
    private JButton saveButton;  // 保存按钮

    public ImageRecognitionUI() {
        setTitle("文字图像识别界面");
        setSize(400, 400);
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setLayout(new FlowLayout());

        imageLabel = new JLabel();  // 创建图像标签
        imageLabel.setPreferredSize(new Dimension(300, 200));  // 设置标签的首选大小
        add(imageLabel);  // 添加到窗口中

        clearButton = new JButton("清空");  // 创建清空按钮
        clearButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                imageLabel.setIcon(null);  // 清空图像标签的图像
                textArea.setText("");  // 清空文本区域的文本
            }
        });
        add(clearButton);  // 添加清空按钮到窗口中

        uploadButton = new JButton("上传图片");  // 创建上传按钮
        uploadButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                JFileChooser fileChooser = new JFileChooser();  // 创建文件选择器
                FileNameExtensionFilter filter = new FileNameExtensionFilter("图像文件", "jpg", "jpeg", "png", "gif");  // 设置过滤器
                fileChooser.setFileFilter(filter);

                int option = fileChooser.showOpenDialog(ImageRecognitionUI.this);  // 显示打开文件对话框
                if (option == JFileChooser.APPROVE_OPTION) {
                    File selectedFile = fileChooser.getSelectedFile();  // 获取选择的文件
                    // 进行图像处理
                    Image processedImage = processImage(selectedFile);
                    // 设置图片标签中的图像
                    imageLabel.setIcon(new ImageIcon(processedImage));  // 设置图像标签的图像为处理后的图像
                    // 进行文字识别
                    String result = recognizeText(selectedFile);
                    textArea.setText(result);  // 设置文本区域的文本为识别结果
                }
            }
        });
        add(uploadButton);  // 添加上传按钮到窗口中

        textArea = new JTextArea(10, 30);  // 创建文本区域
        JScrollPane scrollPane = new JScrollPane(textArea);  // 创建滚动面板
        add(scrollPane);  // 添加滚动面板到窗口中

        saveButton = new JButton("保存结果");  // 创建保存按钮
        saveButton.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e) {
                JFileChooser fileChooser = new JFileChooser();  // 创建文件选择器
                FileNameExtensionFilter filter = new FileNameExtensionFilter("文本文件", "txt");  // 设置过滤器
                fileChooser.setFileFilter(filter);

                int option = fileChooser.showSaveDialog(ImageRecognitionUI.this);  // 显示保存文件对话框
                if (option == JFileChooser.APPROVE_OPTION) {
                    File selectedFile = fileChooser.getSelectedFile();  // 获取选择的文件

                    try {
                        FileWriter fileWriter = new FileWriter(selectedFile);  // 创建文件写入器
                        fileWriter.write(textArea.getText());  // 将文本区域的文本写入文件
                        fileWriter.close();  // 关闭文件写入器
                        JOptionPane.showMessageDialog(ImageRecognitionUI.this, "结果已保存");  // 显示保存成功对话框
                    } catch (IOException ex) {
                        ex.printStackTrace();
                    }
                }
            }
        });
        add(saveButton);  // 添加保存按钮到窗口中

        // 截屏识别按钮
        JButton screenshotButton = new JButton("截屏识别");
        screenshotButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                // 截屏并进行图像识别
                captureScreen();
            }
        });
        add(screenshotButton);
        pack();
        setVisible(true);
    }

    /**
     * 截屏识别
     */
    private void captureScreen() {
        try {
            // 创建一个Robot对象用于屏幕截图
            Robot robot = new Robot();

            // 获取屏幕尺寸
            Dimension screenSize = Toolkit.getDefaultToolkit().getScreenSize();

            // 显示一个具有全屏大小的透明窗口
            JFrame frame = new JFrame();
            frame.setUndecorated(true);
            frame.setOpacity(0.5f);
            frame.setSize(screenSize);
            frame.setLocationRelativeTo(null);
            frame.setCursor(Cursor.getPredefinedCursor(Cursor.CROSSHAIR_CURSOR));
            frame.addMouseListener(new MouseAdapter() {
                private int startX, startY;

                @Override
                public void mousePressed(MouseEvent e) {
                    startX = e.getX();
                    startY = e.getY();
                }

                @Override
                public void mouseReleased(MouseEvent e) {
                    int endX = e.getX();
                    int endY = e.getY();

                    // 计算选取区域的起点和终点坐标
                    int x = Math.min(startX, endX);
                    int y = Math.min(startY, endY);
                    int width = Math.abs(endX - startX);
                    int height = Math.abs(endY - startY);
                    // 进行选区截图
                    BufferedImage screenCapture = robot.createScreenCapture(new Rectangle(x, y, width, height));
                    // 显示选取的图像
                    imageLabel.setIcon(new ImageIcon(screenCapture));
                    //保存文件
                    File screenshotFile = new File(screenshotPath + UUID.randomUUID().toString() + ".png");
                    try {
                        ImageIO.write(screenCapture, "png", screenshotFile);
                        //调用文字识别
                        String recognizeText = recognizeText(screenshotFile);
                        textArea.setText(recognizeText);
                    } catch (IOException ex) {
                        ex.printStackTrace();
                    }
                    frame.dispose();
                }
            });
            frame.setVisible(true);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 图像处理
     */
    private Image processImage(File file) {
        // 调整大小和裁剪
        ImageIcon imageIcon = new ImageIcon(file.getPath());  // 创建图像图标
        return imageIcon.getImage().getScaledInstance(300, 200, Image.SCALE_DEFAULT);
    }

    /**
     * 文字识别
     */
    private String recognizeText(File file) {
        ITesseract tess = new Tesseract();  // 创建Tesseract对象
        tess.setDatapath(tessDataPath);  // 设置Tesseract的数据目录，需要下载语言数据包
        tess.setOcrEngineMode(ITessAPI.TessOcrEngineMode.OEM_DEFAULT);  // 设置OCR引擎模式
        tess.setLanguage("chi_sim+eng");  // 设置语言

        try {
            return tess.doOCR(file);  // 进行OCR识别
        } catch (TesseractException ex) {
            ex.printStackTrace();
            return "文字识别失败";
        }
    }

    public static void main(String[] args) {
        SwingUtilities.invokeLater(new Runnable() {
            public void run() {
                new ImageRecognitionUI();  // 创建图像识别界面对象
            }
        });
    }

}